import os,sys
from torchvision import models
from AI.models.efficientnet import EfficientNet

# 以脚本执行时允许相对导入
if __name__ == "__main__" and __package__ is None:
    sys.path.insert(0, os.path.join(os.getcwd(), '..', '..', '..'))
    import AI  # noqa: F401
    __package__ = "AI"

from AI.models.classifier import  Classifier

if args.pretrained:
    print("Model have pretrained model")
#     model = models.resnet50(pretrained=True)
#     model.fc = nn.Sequential(nn.Linear(2048,512),nn.Linear(512,10),)
#     model = models.resnet18(pretrained=True)
#     model.fc = nn.Sequential(
#         nn.Linear(512,10),)
#     model = EfficientNet.from_pretrained('efficientnet-b0') # 加载预训练模型
else:
    print("Model isn't have pretrained model")
    model = EfficientNet.from_name('efficientnet-b0') # 不加载预训练模型
    model = Classifier()
#     model = models.resnet50(pretrained=False)
#     model.fc = nn.Sequential(nn.Linear(2048,512),nn.Linear(512,10),)
#     model = models.resnet18(pretrained=False)
#     model.fc = nn.Sequential(
#         nn.Linear(512,10),)
